﻿#region LicenseHeader

// Copyright 2012 The Trustees of Indiana University.  All rights reserved.
// 
// The Indiana University licenses this file to you under the Apache License, 
// Version 2.0 (the "License"); you may not use this file except in compliance 
// with the License.  You may obtain a copy of the License at 
// 
// http://www.apache.org/licenses/LICENSE-2.0
// 
// Unless required by applicable law or agreed to in writing, software 
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
//  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  
// See the License for the specific language governing permissions and 
// limitations under the License.
// 
// @author Thilina Gunarathne (tgunarat@indiana.edu)

#endregion

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Threading;
using AzureMRCore.DataModel;
using Microsoft.WindowsAzure;
using Microsoft.WindowsAzure.StorageClient;

namespace AzureMRCore.MapRed
{
    public class MapTaskContext
    {
        private readonly CloudTableClient _cloudTableClient;
        private readonly CloudStorageAccount _storageAccount;
        private readonly string _tableBaseAddress;

        public MapTaskContext(Dictionary<string, object> mapParams, CloudStorageAccount storageAccount)
        {
            AppName = (string) mapParams[Constants.MR_APP_NAME];
            JobID = (string) mapParams[Constants.JOBID];
            NumReduceTasks = Int32.Parse((string) mapParams[Constants.NUM_REDS]);
            MapID = (string) mapParams[Constants.MAPID];
            TaskType = (string) mapParams[Constants.TASK_TYPE];
            IsBlobType = TaskType.Equals(Constants.BLOB_TYPE);
            if (IsBlobType)
            {
                BlobUrl = (string) mapParams[Constants.BLOB_URI];
            }
            ProgramParams = (string) mapParams[Constants.PROGRAM_PARAMS];
            OutDir = (string) mapParams[Constants.OUT_DIR];
            if (mapParams.ContainsKey(Constants.ITERATION)) Iteration = (string) mapParams[Constants.ITERATION];
            if (TaskType.Equals(Constants.CUSTOM_TYPE))
            {
                IsCustomType = true;
                CustomKey = (string) mapParams[Constants.CUSTOM_KEY];
                CustomValue = (string) mapParams[Constants.CUSTOM_VALUE];
            }

            _storageAccount = storageAccount;
            _tableBaseAddress = _storageAccount.TableEndpoint.ToString();
            _cloudTableClient = _storageAccount.CreateCloudTableClient();

            MapDataContext = new MapMetaDataContext(_tableBaseAddress,
                                                    storageAccount.Credentials,
                                                    DataModelUtils.GetMapTableName(JobID));
            MapTask = GetMapTaskEntity((string) mapParams[Constants.MAPID], MapDataContext);
            IsPrimitiveBCast = false;
        }

        public MapTaskContext(CloudStorageAccount storageAccount)
        {
            _storageAccount = storageAccount;
            _tableBaseAddress = _storageAccount.TableEndpoint.ToString();
            _cloudTableClient = _storageAccount.CreateCloudTableClient();
            IsPrimitiveBCast = false;
        }

        public MapDataModel MapTask { get; set; }
        public MapMetaDataContext MapDataContext { get; set; }
        public int NumReduceTasks { get; set; }
        public string AppName { get; set; }
        public string JobID { get; set; }
        public string BlobUrl { get; set; }
        public string TaskType { get; set; }
        public string ProgramParams { get; set; }
        public string OutDir { get; set; }
        public string MapID { get; set; }
        public string Iteration { get; set; }
        public string CustomKey { get; set; }
        public string CustomValue { get; set; }
        public int Attempt { get; set; }
        public bool IsBlobType { get; set; }
        public bool IsCustomType { get; set; }
        // set by direct TCP transfer
        public ReduceMetaDataContext ReduceMetaDataContext { get; set; }
        // we can get rid of this when we move all BCast to receiver daemon
        public bool IsPrimitiveBCast { get; set; }

        private MapDataModel GetMapTaskEntity(string taskID, MapMetaDataContext mapDataContext)
        {
            int retryCount = 0;
            mapTableEntryRetryLabel:
            IQueryable<MapDataModel> mapTaskEnum =
                mapDataContext.MapTasks.Where(p => (p.RowKey == taskID) & p.PartitionKey == Iteration);

            MapDataModel mapTask = mapTaskEnum.FirstOrDefault();
            if (mapTask == null)
            {
                if (retryCount < 200)
                {
                    retryCount++;
                    Thread.Sleep(retryCount*3);
                    goto mapTableEntryRetryLabel;
                }
                throw new Exception("No table entry for the map task. taskID:" + taskID + " JobID:" + JobID +
                                    " iteration:" + Iteration + " retries:" + retryCount);
            }
            return mapTask;
        }

        /*
         * This can either be reduce task table or broadcast data table
         */

        internal ReduceInputDataContext CreateIntermediateDataTable(int id, bool isBroadcast)
        {
            ReduceInputDataContext context;
            int retry = 0;
            // create the table here, which saves us time in the client
            string taskTableName = isBroadcast
                                       ? DataModelUtils.GetBroadcastTaskTableName(JobID, Iteration, AppName)
                                       : DataModelUtils.GetReduceTaskTableName(JobID, id, AppName);

            reduceTaskTableCreateLabel:
            try
            {
                _cloudTableClient.CreateTableIfNotExist(taskTableName);
                context = new ReduceInputDataContext(_tableBaseAddress, _storageAccount.Credentials,
                                                     taskTableName);
            }

            catch (Exception e)
            {
                if (retry < Constants.NUM_RETRIES)
                {
                    retry++;
                    Thread.Sleep(50);
                    goto reduceTaskTableCreateLabel;
                }
                throw new Exception("reduceTaskTableCreateLabel failed. retry " + retry, e);
            }
            return context;
        }

        internal void RecordTimeAndSave(Stopwatch mapTimer)
        {
            MapTask.MapTaskTime = mapTimer.ElapsedMilliseconds;
            MapDataContext.UpdateObject(MapTask);
            MapDataContext.SaveChangesWithRetries();
        }

        internal void UpdateMapTaskStatus(int mapTaskStatus, byte[] reduceProductCounts)
        {
            switch (Attempt)
            {
                case 1:
                    MapTask.Attempt1Status = mapTaskStatus;
                    break;
                case 2:
                    MapTask.Attempt2Status = mapTaskStatus;
                    break;
                case 3:
                    MapTask.Attempt3Status = mapTaskStatus;
                    break;
            }
            if (mapTaskStatus == MapDataModel.FINISHED)
            {
                MapTask.EndTime = string.Format("{0:G}", DateTime.Now);
                MapTask.TaskStatus = MapDataModel.FINISHED;
                MapTask.ReduceProductCounts = reduceProductCounts;
            }
        }

        internal int SetStateToExecute()
        {
            if (MapTask.Attempt1Status == MapDataModel.SUBMITTED)
            {
                MapTask.Attempt1Status = MapDataModel.EXECUTING;
                MapTask.Attempt1StartTime = string.Format("{0:G}", DateTime.Now);
                Attempt = 1;
            }
            else if (MapTask.Attempt2Status == MapDataModel.SUBMITTED)
            {
                MapTask.Attempt2Status = MapDataModel.EXECUTING;
                MapTask.Attempt2StartTime = string.Format("{0:G}", DateTime.Now);
                Attempt = 2;
            }
            else if (MapTask.Attempt3Status == MapDataModel.SUBMITTED)
            {
                MapTask.Attempt3Status = MapDataModel.EXECUTING;
                MapTask.Attempt3StartTime = string.Format("{0:G}", DateTime.Now);
                Attempt = 3;
            }
            else
            {
                //to do handle job failure
                return -1;
            }
            MapTask.TaskStatus = MapDataModel.EXECUTING;
            MapDataContext.UpdateObject(MapTask);
            MapDataContext.SaveChanges();
            return Attempt;
        }
    }
}